#ifndef ARMOR_TRACKER__TRACKER_HPP
#define ARMOR_TRACKER__TRACKER_HPP

#include "aim_interfaces/msg/armor.hpp"
#include "aim_interfaces/msg/armors.hpp"

#include "extended_kalman_filter.hpp"
#include "rclcpp/time.hpp"

#include <eigen3/Eigen/Dense>
#include <string>

namespace rm_aim
{
/**
 * @brief 装甲板数量类型枚举
 * @note 不同类型目标的装甲板数量不同，用于处理多装甲板切换逻辑
 */
enum class ArmorCount {
  kNormal4 = 0,   // 常规目标（4个装甲板）
  kBalance2 = 1,  // 平衡步兵（2个装甲板）
  kOutpost3 = 2   // 前哨站（3个装甲板）
};

/**
 * @brief 追踪状态枚举
 * @note 定义追踪器的不同工作状态，用于状态机管理
 */
enum class TrackingState {
  kLost = 0,      // 丢失目标
  kDetecting = 1, // 检测中（未稳定追踪）
  kTracking = 2,  // 稳定追踪中
  kTempLost = 3   // 暂时丢失（短时间未检测到）
};

/**
 * @brief 装甲板追踪器类
 * @details 基于扩展卡尔曼滤波器（EKF）实现对装甲板目标的实时追踪，
 *          包含状态机管理、装甲板匹配、EKF初始化与更新等功能
 */
class ArmorTracker
{
public:
  /**
   * @brief 构造函数
   * @param max_match_dist 装甲板匹配最大距离阈值（单位：m）
   * @param max_match_yaw 装甲板匹配最大偏航角差阈值（单位：rad）
   */
  ArmorTracker(double max_match_dist, double max_match_yaw);

  /**
   * @brief 初始化追踪器
   * @param armors_msg 装甲板检测消息（包含多个装甲板信息）
   * @note 从检测消息中选择初始追踪目标，初始化EKF和追踪状态
   */
  void initialize(const aim_interfaces::msg::Armors::SharedPtr & armors_msg);

  /**
   * @brief 更新追踪状态
   * @param armors_msg 最新的装甲板检测消息
   * @note 执行EKF预测、装甲板匹配、EKF更新，并更新追踪状态机
   */
  void updateTracking(const aim_interfaces::msg::Armors::SharedPtr & armors_msg);

  /**
   * @brief 设置EKF参数（过程噪声和观测噪声系数）
   * @param sigma_q_xyz 位置/速度过程噪声方差系数
   * @param sigma_q_yaw 偏航角/角速度过程噪声方差系数
   * @param sigma_q_r 装甲板半径过程噪声方差系数
   * @param r_xyz_factor 位置观测噪声系数
   * @param r_yaw 偏航角观测噪声值
   */
  void setEKFParams(
    double sigma_q_xyz, double sigma_q_yaw, double sigma_q_r,
    double r_xyz_factor, double r_yaw);

  /**
   * @brief 获取当前追踪状态
   * @return 追踪状态（枚举值）
   */
  TrackingState getTrackingState() const { return tracking_state_; }
  
  /**
   * @brief 设置当前追踪状态
   * @param state 追踪状态（枚举值）
   */
  void setTrackingState(TrackingState state) { tracking_state_ = state; }

  /**
   * @brief 记录chi-square检验结果
   */
  void logChiSquareValidation();

  // 公开成员（供节点获取追踪结果）
  std::string tracked_armor_id;          // 当前追踪的装甲板ID
  aim_interfaces::msg::Armor tracked_armor;  // 当前匹配的装甲板
  Eigen::VectorXd target_state;          // 目标状态向量（9维）
  Eigen::VectorXd measurement;           // 观测值向量（4维：x,y,z,yaw）
  double info_position_diff;             // 调试信息：位置偏差
  double info_yaw_diff;                  // 调试信息：偏航角偏差
  ArmorCount tracked_armor_count;        // 当前目标的装甲板数量类型
  double dz;                             // 装甲板高度差（多装甲板切换用）
  double another_r;                      // 备用半径（多装甲板切换用）
  int tracking_thres;                    // 检测转追踪的连续匹配阈值
  int lost_thres;                        // 暂时丢失转丢失的计数阈值

private:
  /**
   * @brief 初始化EKF
   * @param armor 初始装甲板信息
   * @note 根据初始装甲板位置和角度，计算目标中心初始状态，初始化EKF
   */
  void initEKF(const aim_interfaces::msg::Armor & armor);

  /**
   * @brief 处理装甲板跳变（目标旋转导致的装甲板切换）
   * @param current_armor 当前检测到的同ID装甲板
   * @note 修正EKF状态，切换半径和高度参数，必要时重置EKF
   */
  void handleArmorJump(const aim_interfaces::msg::Armor & current_armor);

  /**
   * @brief 更新目标的装甲板数量类型
   * @param armor 当前装甲板信息
   * @note 根据装甲板ID和类型判断目标属于哪种装甲板数量类型
   */
  void updateArmorCount(const aim_interfaces::msg::Armor & armor);

  /**
   * @brief 将四元数转换为偏航角（处理角度连续性）
   * @param q 四元数姿态
   * @return 连续的偏航角（单位：rad）
   * @note 解决偏航角在-pi~pi之间的跳变问题，输出连续角度
   */
  double quaternionToYaw(const geometry_msgs::msg::Quaternion & q);

  /**
   * @brief 从目标状态计算装甲板位置
   * @param state 目标状态向量
   * @return 装甲板位置（x,y,z）
   * @note 根据目标中心位置、偏航角和半径，反推装甲板位置
   */
  Eigen::Vector3d getArmorPosFromState(const Eigen::VectorXd & state);

  // 核心成员
  TrackingState tracking_state_;         // 当前追踪状态
  tools::ExtendedKalmanFilter ekf_;      // 扩展卡尔曼滤波器实例

  // 时间与采样参数
  rclcpp::Time last_update_stamp_;       // 上一次更新时间戳
  double sampling_time_;                 // 采样时间间隔（单位：s）

  // 匹配阈值参数
  const double max_match_distance_;      // 最大匹配距离阈值（m）
  const double max_match_yaw_diff_;      // 最大匹配偏航角差阈值（rad）

  // 计数变量
  int detect_count_;                     // 连续检测计数（用于状态切换）
  int lost_count_;                       // 丢失计数（用于状态切换）
  double last_yaw_;                      // 上一时刻偏航角（用于角度连续性处理）

  // EKF参数（从节点传入）
  double sigma2_q_xyz_;                  // 位置/速度过程噪声方差系数
  double sigma2_q_yaw_;                  // 偏航角/角速度过程噪声方差系数
  double sigma2_q_r_;                    // 半径过程噪声方差系数
  double r_xyz_factor_;                  // 位置观测噪声系数
  double r_yaw_;                         // 偏航角观测噪声值

  // EKF初始协方差矩阵（固定参数）
  const Eigen::MatrixXd init_covariance_ = []() {
    Eigen::MatrixXd P0(9, 9);
    P0.setIdentity();
    // 初始协方差：位置(x,y,z)方差0.1，速度方差0.01，角度方差0.05，半径方差0.001
    P0.diagonal() << 0.1, 0.01, 0.1, 0.01, 0.1, 0.01, 0.05, 0.005, 0.001;
    return P0;
  }();
};

}  // namespace rm_aim

#endif  // ARMOR_TRACKER__TRACKER_HPP